import os

def replace_printline_in_c(txt_dir, c_dir):
    # 获取txt目录下的所有txt文件
    txt_files = [f for f in os.listdir(txt_dir) if f.endswith('.txt')]
    # 获取c目录下的所有c文件
    c_files = [f for f in os.listdir(c_dir) if (f.endswith('.c') or f.endswith('.cpp'))]

    # 如果没有txt文件或c文件，给出警告并返回
    if not txt_files or not c_files:
        print("Error: No .txt or .c files found.")
        return

    # 计算txt文件的数量
    txt_file_count = len(txt_files)
    
    # 循环遍历所有.c文件并用txt文件内容进行替换
    for i, c_filename in enumerate(c_files):
        # 循环使用txt文件，确保用完后继续循环使用
        txt_filename = txt_files[i % txt_file_count]  # 循环访问txt文件

        txt_file_path = os.path.join(txt_dir, txt_filename)
        c_file_path = os.path.join(c_dir, c_filename)

        # 读取txt文件内容
        with open(txt_file_path, 'r', encoding='utf-8') as txt_file:
            txt_content = txt_file.read().strip()  # 读取并去掉首尾空格

        # 读取c文件内容
        with open(c_file_path, 'r', encoding='utf-8') as c_file:
            c_content = c_file.readlines()

        # 查找并替换printLine("Calling good()..."); 和 printLine("Calling bad()..."); 行
        modified_content = []
        i = 0
        while i < len(c_content):
            line = c_content[i]
            if 'printLine("Calling good()...");' in line or 'printLine("Calling bad()...");' in line:
                # 完全替换printLine语句
                modified_content.append(f'{txt_content}')  # 直接用txt文件中的内容替换
                # 查找最近的以 `);` 结尾的语句，修改它
                j = i + 1
                while j < len(c_content):
                    next_line = c_content[j].strip()
                    if next_line.endswith(');'):
                        # 找到以 `);` 结尾的语句后，修改此行
                        c_content[j] = next_line + ' }\n'  # 在当前行末尾添加 `}`
                        break
                    j += 1
            else:
                modified_content.append(line)
            i += 1

        # 将修改后的内容写回文件
        with open(c_file_path, 'w', encoding='utf-8') as c_file:
            c_file.writelines(modified_content)

        print(f'Updated {c_filename} with content from {txt_filename}')

# 输入txt和c目录路径
txt_directory = '/home/kali/Desktop/angr/deadcode'  # 这里替换为你的txt文件目录
c_directory = '/home/kali/Desktop/angr/2017-10-01-juliet-test-suite-for-c-cplusplus-v1-3/C/testcases/CWE121_Stack_Based_Buffer_Overflow/s07'      # 这里替换为你的c文件目录

# 调用函数
replace_printline_in_c(txt_directory, c_directory)
